// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <nebula/array/builder_dict.h>

#include <cstdint>
#include <limits>
#include <memory>
#include <string_view>
#include <type_traits>
#include <utility>
#include <vector>

#include <nebula/core/array.h>
#include <nebula/core/buffer.h>

#include <turbo/utility/status.h>
#include <nebula/types/type.h>
#include <nebula/types/type_traits.h>
#include <turbo/base/checked_cast.h>
#include <nebula/util/hashing.h>
#include <turbo/log/logging.h>

namespace nebula {
    namespace internal {

        template<typename T, typename Enable = void>
        struct DictionaryTraits {
            using MemoTableType = void;
        };

    }  // namespace internal

    template<typename T, typename Out = void>
    using enable_if_memoize = enable_if_t<
            !std::is_same<typename internal::DictionaryTraits<T>::MemoTableType, void>::value,
            Out>;

    template<typename T, typename Out = void>
    using enable_if_no_memoize = enable_if_t<
            std::is_same<typename internal::DictionaryTraits<T>::MemoTableType, void>::value,
            Out>;

    namespace internal {

        template<>
        struct DictionaryTraits<BooleanType> {
            using T = BooleanType;
            using MemoTableType = typename HashTraits<T>::MemoTableType;

            static turbo::Result<std::shared_ptr<ArrayData>> GetDictionaryArrayData(
                    MemoryPool *pool, const std::shared_ptr<DataType> &type,
                    const MemoTableType &memo_table, int64_t start_offset) {
                if (start_offset < 0) {
                    return turbo::invalid_argument_error("invalid start_offset ", start_offset);
                }

                BooleanBuilder builder(pool);
                const auto &bool_values = memo_table.values();
                const auto null_index = memo_table.GetNull();

                // Will iterate up to 3 times.
                for (int64_t i = start_offset; i < memo_table.size(); i++) {
                    TURBO_RETURN_NOT_OK(i == null_index ? builder.append_null()
                                                        : builder.append(bool_values[i]));
                }

                std::shared_ptr<ArrayData> out;
                TURBO_RETURN_NOT_OK(builder.finish_internal(&out));
                return out;
            }
        };  // namespace internal

        template<typename T>
        struct DictionaryTraits<T, enable_if_has_c_type<T>> {
            using c_type = typename T::c_type;
            using MemoTableType = typename HashTraits<T>::MemoTableType;

            static turbo::Result<std::shared_ptr<ArrayData>> GetDictionaryArrayData(
                    MemoryPool *pool, const std::shared_ptr<DataType> &type,
                    const MemoTableType &memo_table, int64_t start_offset) {
                auto dict_length = static_cast<int64_t>(memo_table.size()) - start_offset;
                // This makes a copy, but we assume a dictionary array is usually small
                // compared to the size of the dictionary-using array.
                // (also, copying the dictionary values is cheap compared to the cost
                //  of building the memo table)
                TURBO_MOVE_OR_RAISE(
                        std::shared_ptr<Buffer> dict_buffer,
                        allocate_buffer(TypeTraits<T>::bytes_required(dict_length), pool));
                memo_table.CopyValues(static_cast<int32_t>(start_offset),
                                      reinterpret_cast<c_type *>(dict_buffer->mutable_data()));

                int64_t null_count = 0;
                std::shared_ptr<Buffer> null_bitmap = nullptr;
                TURBO_RETURN_NOT_OK(
                        ComputeNullBitmap(pool, memo_table, start_offset, &null_count, &null_bitmap));

                return ArrayData::create(type, dict_length, {null_bitmap, dict_buffer}, null_count);
            }
        };

        template<typename T>
        struct DictionaryTraits<T, enable_if_base_binary<T>> {
            using MemoTableType = typename HashTraits<T>::MemoTableType;

            static turbo::Result<std::shared_ptr<ArrayData>> GetDictionaryArrayData(
                    MemoryPool *pool, const std::shared_ptr<DataType> &type,
                    const MemoTableType &memo_table, int64_t start_offset) {
                using offset_type = typename T::offset_type;

                // create the offsets buffer
                auto dict_length = static_cast<int64_t>(memo_table.size() - start_offset);
                TURBO_MOVE_OR_RAISE(auto dict_offsets,
                                    allocate_buffer(sizeof(offset_type) * (dict_length + 1), pool));
                auto raw_offsets = reinterpret_cast<offset_type *>(dict_offsets->mutable_data());
                memo_table.CopyOffsets(static_cast<int32_t>(start_offset), raw_offsets);

                // create the data buffer
                auto values_size = memo_table.values_size();
                TURBO_MOVE_OR_RAISE(auto dict_data, allocate_buffer(values_size, pool));
                if (values_size > 0) {
                    memo_table.CopyValues(static_cast<int32_t>(start_offset), dict_data->size(),
                                          dict_data->mutable_data());
                }

                int64_t null_count = 0;
                std::shared_ptr<Buffer> null_bitmap = nullptr;
                TURBO_RETURN_NOT_OK(
                        ComputeNullBitmap(pool, memo_table, start_offset, &null_count, &null_bitmap));

                return ArrayData::create(type, dict_length,
                                         {null_bitmap, std::move(dict_offsets), std::move(dict_data)},
                                         null_count);
            }
        };

        template<typename T>
        struct DictionaryTraits<T, enable_if_binary_view_like<T>> {
            using MemoTableType = typename HashTraits<T>::MemoTableType;

            static_assert(std::is_same_v<MemoTableType, BinaryMemoTable<BinaryBuilder>>);

            // Instead of defining a custom memo table for StringView we reuse BinaryType's,
            // then convert to views when we copy data out of the memo table.
            static turbo::Result<std::shared_ptr<ArrayData>> GetDictionaryArrayData(
                    MemoryPool *pool, const std::shared_ptr<DataType> &type,
                    const MemoTableType &memo_table, int64_t start_offset) {
                        DKCHECK(type->id() == Type::STRING_VIEW || type->id() == Type::BINARY_VIEW);

                BinaryViewBuilder builder(pool);
                TURBO_RETURN_NOT_OK(builder.resize(memo_table.size() - start_offset));
                TURBO_RETURN_NOT_OK(builder.reserve_data(memo_table.values_size()));
                memo_table.VisitValues(static_cast<int32_t>(start_offset),
                                       [&](std::string_view s) { builder.unsafe_append(s); });

                std::shared_ptr<ArrayData> out;
                TURBO_RETURN_NOT_OK(builder.finish_internal(&out));
                out->type = type;
                return out;
            }
        };

        template<typename T>
        struct DictionaryTraits<T, enable_if_fixed_size_binary<T>> {
            using MemoTableType = typename HashTraits<T>::MemoTableType;

            static turbo::Result<std::shared_ptr<ArrayData>> GetDictionaryArrayData(
                    MemoryPool *pool, const std::shared_ptr<DataType> &type,
                    const MemoTableType &memo_table, int64_t start_offset) {
                const T &concrete_type = turbo::checked_cast<const T &>(*type);

                // create the data buffer
                auto dict_length = static_cast<int64_t>(memo_table.size() - start_offset);
                auto width_length = concrete_type.byte_width();
                auto data_length = dict_length * width_length;
                TURBO_MOVE_OR_RAISE(auto dict_data, allocate_buffer(data_length, pool));
                auto data = dict_data->mutable_data();

                memo_table.CopyFixedWidthValues(static_cast<int32_t>(start_offset), width_length,
                                                data_length, data);

                int64_t null_count = 0;
                std::shared_ptr<Buffer> null_bitmap = nullptr;
                TURBO_RETURN_NOT_OK(
                        ComputeNullBitmap(pool, memo_table, start_offset, &null_count, &null_bitmap));

                return ArrayData::create(type, dict_length, {null_bitmap, std::move(dict_data)},
                                         null_count);
            }
        };

    }  // namespace internal
}  // namespace nebula
